# INSERT_YOUR_CODE
import argparse
import os

parser = argparse.ArgumentParser()
parser.add_argument('-c', '--cuda', type=str, default=None, help='CUDA_VISIBLE_DEVICES setting')
args, unknown = parser.parse_known_args()

if args.cuda is not None:
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda



from PIL import Image
from util_pose import get_pose_predictor
from util_parse import get_parse_predictor
from util_mask_tryon import tryon_pipeline
from tqdm import tqdm

pose_predictor = get_pose_predictor()
parse_predictor = get_parse_predictor()

data_root = "/mnt/nas/shengjie/datasets/DressCode_1024"
# clo_types = ["dresses", "upper", "lower"]
sub_dirs = [
    "cloth", "cloth_align_mask-bytedance", "cloth_mask", "densepose", "parse-bytedance",
    "cloth_align", "cloth_align_parse-bytedance", "cloth_parse", "image", "pose_25"
]
get_save_dir = lambda clo_type: f'/mnt/nas/shengjie/datasets/tryon_{clo_type}_unpaired'

unpaired_list_path = '/mnt/nas/shengjie/datasets/DressCode_1024/test_pairs_unpaired_230729.txt'

'''
(flux2) ... head /mnt/nas/shengjie/datasets/DressCode_1024/test_pairs_unpaired_230729.txt 
048392_0.png 049114_1.png upper
048408_0.png 048433_1.png upper
048409_0.png 049910_1.png upper

根据第三列的分类 以 (human_name, clo_name) 即 (048392_0, 049114_1)的形式存入 paired_names_{category}
构造 clo_dicts = { {category} : paired_names_{category}  }

(Pdb) clo_dicts.keys()
dict_keys(['upper', 'lower', 'dresses'])
(Pdb) len(clo_dicts['upper'])
1551
(Pdb) len(clo_dicts['lower'])
1210
(Pdb) len(clo_dicts['dresses'])
1527
'''
# 读取 unpaired list 文件，按照类型组织成字典
clo_dicts = {}
with open(unpaired_list_path, "r") as f:
    for line in f:
        items = line.strip().split()
        if len(items) != 3:
            continue  # 跳过格式错误的行
        human_png, clo_png, category = items
        # 忽略扩展名，得到纯stem（不带.png）
        human_name = os.path.splitext(human_png)[0]
        clo_name = os.path.splitext(clo_png)[0]
        pair = (human_name, clo_name)
        if category not in clo_dicts:
            clo_dicts[category] = []
        clo_dicts[category].append(pair)

for clo_type, unpaired_names in clo_dicts.items():
    main_dir = os.path.join(data_root, clo_type)

    human_dir = os.path.join(main_dir, 'image')
    cloth_dir = os.path.join(main_dir, 'cloth')
    cloth_mask_dir = os.path.join(main_dir, 'cloth_mask')


    # INSERT_YOUR_CODE
    # 对于每个clo图片名，找到对应的human图片名，确保两个路径都存在
    for human_name, clo_name in tqdm(unpaired_names):
        # clo_name: e.g. 000000_1.jpg
        human_name = human_name + '.jpg'
        clo_mask_name = clo_name + '.png'
        clo_name = clo_name + '.jpg'
        clo_path = os.path.join(cloth_dir, clo_name)
        clo_mask_path = os.path.join(cloth_mask_dir, clo_mask_name)
        human_path = os.path.join(human_dir, human_name)
        if not os.path.isfile(clo_path):
            print(f"Cloth image {clo_path} does not exist, skipping.")
            continue
        if not os.path.isfile(clo_mask_path):
            print(f"Cloth mask image {clo_mask_path} does not exist, skipping.")
            continue
        if not os.path.isfile(human_path):
            print(f"Human image {human_path} does not exist, skipping.")
            continue

        # 可选: 运行tryon_pipeline做一次推理和可视化
        try:
            result = tryon_pipeline(
                img_path=human_path,
                clo_path=clo_path,
                clo_type=clo_type if clo_type!='dresses' else 'full',
                clo_mask_path=clo_mask_path,
                pose_predictor=pose_predictor,
                parse_predictor=parse_predictor
            )
            # 保存可视化结果
            save_dir = get_save_dir(clo_type)
            os.makedirs(save_dir, exist_ok=True)
            out_name = os.path.splitext(clo_name)[0] + "_tryon.jpg"
            out_path = os.path.join(save_dir, out_name)
            result["final_img"].save(out_path)
            print(f"Saved tryon visualization to {out_path}")
            print()
        except Exception as e:
            print(f"Failed tryon for {clo_name}: {e}")

# 后续 还需要在提示词中添加 skin 颜色


'''
cloth RGB 
cloth_align_mask-bytedance  L array([  0, 255], dtype=uint8) cloth's unwarp mask == cloth_mask 
cloth_mask                  L array([  0, 255], dtype=uint8) cloth's unwarp mask
densepose       L       array([ 0,  2,  3,  4,  5,  6,  7,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
                                19, 20, 21, 22, 23, 24], dtype=uint8) human's densepose
parse-bytedance
cloth_align RGB == cloth
cloth_align_parse-bytedance  L   array([0, 5], dtype=uint8) cloth's parse
cloth_parse L == cloth_align_parse-bytedance
image RGB human
pose_25 .npy 数据 pose points 可视化为pose图  
        pose 我重新拿模型预测吧
'''